# test_fixed.py
import asyncio
import websockets
import json
import uuid
from datetime import datetime
import time
import requests
import logging

# 配置测试日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%H:%M:%S'
)
logger = logging.getLogger(__name__)


class WebSocketServerTester:
    """修复的 WebSocket 服务器测试类"""

    def __init__(self, base_url="http://localhost:8000", ws_url="ws://localhost:8000"):
        self.base_url = base_url
        self.ws_url = ws_url
        self.connected_devices = {}
        self.timeout = 10  # 增加超时时间

    async def test_device_registration(self, device_count=2):
        """测试设备注册功能 - 修复版本"""
        logger.info(f"🧪 测试设备注册 - 准备注册 {device_count} 个设备")

        success_count = 0
        for i in range(device_count):
            device_id = f"test_device_{i + 1}"
            if await self._register_single_device(device_id):
                success_count += 1
            await asyncio.sleep(0.5)  # 添加延迟避免同时连接

        logger.info(f"✅ 设备注册测试完成: {success_count}/{device_count} 成功")
        return success_count > 0

    async def _register_single_device(self, device_id: str):
        """注册单个设备 - 修复版本"""
        try:
            # 增加连接超时时间
            websocket = await asyncio.wait_for(
                websockets.connect(f"{self.ws_url}/ws/device/{device_id}"),
                timeout=self.timeout
            )

            # 发送注册消息
            register_message = {
                "type": "device_register",
                "data": {
                    "deviceId": device_id,
                    "versionName": "v1.0.0",
                    "url": f"ws://client/{device_id}",
                    "code": f"CODE_{device_id.upper()}",
                    "capabilities": ["temperature", "humidity", "pressure"]
                },
                "timestamp": datetime.now().isoformat()
            }

            await websocket.send(json.dumps(register_message))
            logger.info(f"📤 已发送注册消息给设备 {device_id}")

            # 等待响应
            try:
                response = await asyncio.wait_for(websocket.recv(), timeout=self.timeout)
                response_data = json.loads(response)
                logger.info(f"📥 收到设备 {device_id} 的响应: {response_data.get('type')}")

                if response_data.get("type") == "status_update":
                    if response_data.get("data", {}).get("status") == "registered":
                        self.connected_devices[device_id] = websocket
                        logger.info(f"✅ 设备 {device_id} 注册成功")
                        return True
                    else:
                        logger.error(f"❌ 设备 {device_id} 注册状态异常: {response_data}")
                        await websocket.close()
                        return False
                else:
                    logger.error(f"❌ 设备 {device_id} 收到意外响应: {response_data}")
                    await websocket.close()
                    return False

            except asyncio.TimeoutError:
                logger.error(f"⏰ 设备 {device_id} 注册响应超时")
                await websocket.close()
                return False

        except asyncio.TimeoutError:
            logger.error(f"⏰ 设备 {device_id} 连接超时")
            return False
        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 连接失败: {str(e)}")
            return False

    async def test_heartbeat(self, device_id: str = "test_device_1"):
        """测试心跳功能 - 修复版本"""
        logger.info(f"🧪 测试心跳功能 - 设备: {device_id}")

        try:
            if device_id not in self.connected_devices:
                logger.warning(f"⚠️  设备 {device_id} 未连接，先进行注册")
                if not await self._register_single_device(device_id):
                    return False

            websocket = self.connected_devices[device_id]

            # 发送心跳消息
            heartbeat_message = {
                "type": "heartbeat",
                "data": {
                    "status": "healthy",
                    "cpu_usage": 45.2,
                    "memory_usage": 67.8,
                    "uptime": 3600
                },
                "timestamp": datetime.now().isoformat()
            }

            await websocket.send(json.dumps(heartbeat_message))
            logger.info(f"📤 已发送心跳消息给设备 {device_id}")

            # 等待心跳响应
            response = await asyncio.wait_for(websocket.recv(), timeout=self.timeout)
            response_data = json.loads(response)

            if response_data.get("type") == "heartbeat":
                logger.info(f"✅ 设备 {device_id} 心跳测试成功")
                return True
            else:
                logger.error(f"❌ 设备 {device_id} 心跳测试失败: {response_data}")
                return False

        except asyncio.TimeoutError:
            logger.error(f"⏰ 设备 {device_id} 心跳响应超时")
            return False
        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 心跳测试异常: {str(e)}")
            return False

    async def test_status_update(self, device_id: str = "test_device_1"):
        """测试状态更新功能 - 修复版本"""
        logger.info(f"🧪 测试状态更新 - 设备: {device_id}")

        try:
            if device_id not in self.connected_devices:
                if not await self._register_single_device(device_id):
                    return False

            websocket = self.connected_devices[device_id]

            # 发送状态更新消息
            status_message = {
                "type": "status_update",
                "data": {
                    "temperature": 25.5,
                    "humidity": 65.2,
                    "pressure": 1013.25,
                    "battery_level": 85,
                    "network_strength": 4
                },
                "timestamp": datetime.now().isoformat()
            }

            await websocket.send(json.dumps(status_message))
            logger.info(f"📤 已发送状态更新给设备 {device_id}")

            # 状态更新不需要响应，等待一下确保消息发送完成
            await asyncio.sleep(1)
            logger.info(f"✅ 设备 {device_id} 状态更新测试成功")
            return True

        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 状态更新测试异常: {str(e)}")
            return False

    async def test_data_report(self, device_id: str = "test_device_1"):
        """测试数据报告功能 - 修复版本"""
        logger.info(f"🧪 测试数据报告 - 设备: {device_id}")

        try:
            if device_id not in self.connected_devices:
                if not await self._register_single_device(device_id):
                    return False

            websocket = self.connected_devices[device_id]

            # 发送数据报告消息
            data_message = {
                "type": "data_report",
                "data": {
                    "data_points": [
                        {
                            "timestamp": datetime.now().isoformat(),
                            "value": 25.1,
                            "sensor": "temperature",
                            "unit": "celsius"
                        },
                        {
                            "timestamp": datetime.now().isoformat(),
                            "value": 65.3,
                            "sensor": "humidity",
                            "unit": "percent"
                        }
                    ],
                    "total_count": 2,
                    "sampling_rate": "1s"
                },
                "timestamp": datetime.now().isoformat()
            }

            await websocket.send(json.dumps(data_message))
            logger.info(f"📤 已发送数据报告给设备 {device_id}")

            await asyncio.sleep(1)
            logger.info(f"✅ 设备 {device_id} 数据报告测试成功")
            return True

        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 数据报告测试异常: {str(e)}")
            return False

    async def test_command_receiving(self, device_id: str = "test_device_1"):
        """测试接收服务器命令 - 修复版本"""
        logger.info(f"🧪 测试命令接收 - 设备: {device_id}")
        logger.info("等待服务器发送定时命令（约30秒间隔）...")

        try:
            if device_id not in self.connected_devices:
                if not await self._register_single_device(device_id):
                    return False

            websocket = self.connected_devices[device_id]

            # 设置一个消息接收任务
            async def receive_command():
                try:
                    # 等待服务器发送命令（缩短等待时间用于测试）
                    start_time = time.time()
                    while time.time() - start_time < 35:  # 最多等待35秒
                        try:
                            message = await asyncio.wait_for(websocket.recv(), timeout=5)
                            message_data = json.loads(message)
                            if message_data.get("type") == "command":
                                logger.info(f"✅ 设备 {device_id} 成功接收到命令: {message_data.get('data', {}).get('command')}")
                                return True
                        except asyncio.TimeoutError:
                            continue  # 继续等待

                    logger.warning(f"⚠️  设备 {device_id} 未在35秒内收到命令")
                    return False

                except Exception as e:
                    logger.error(f"❌ 命令接收过程异常: {str(e)}")
                    return False

            result = await receive_command()
            return result

        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 命令接收测试异常: {str(e)}")
            return False

    def test_rest_api(self):
        """测试 REST API 接口 - 修复版本"""
        logger.info("🧪 测试 REST API 接口")

        try:
            # 测试根路径
            response = requests.get(f"{self.base_url}/", timeout=5)
            if response.status_code == 200:
                logger.info("✅ 根路径测试成功")
            else:
                logger.error(f"❌ 根路径测试失败: {response.status_code}")
                return False

            # 测试获取在线设备
            response = requests.get(f"{self.base_url}/devices/online", timeout=5)
            if response.status_code == 200:
                online_devices = response.json()
                logger.info(f"✅ 在线设备查询成功: {len(online_devices)} 个设备在线")
            else:
                logger.error(f"❌ 在线设备查询失败: {response.status_code}")
                return False

            # 测试获取特定设备状态
            response = requests.get(f"{self.base_url}/devices/test_device_1", timeout=5)
            if response.status_code == 200:
                device_status = response.json()
                logger.info(f"✅ 设备状态查询成功: {device_status.get('deviceId')} - 在线: {device_status.get('is_online')}")
            elif response.status_code == 404:
                logger.warning("⚠️  设备状态查询 - 设备不存在（正常情况）")
            else:
                logger.error(f"❌ 设备状态查询失败: {response.status_code}")
                return False

            return True

        except requests.exceptions.Timeout:
            logger.error("⏰ REST API 请求超时")
            return False
        except Exception as e:
            logger.error(f"❌ REST API 测试异常: {str(e)}")
            return False

    async def test_broadcast_functionality(self):
        """测试广播功能 - 修复版本"""
        logger.info("🧪 测试广播功能")

        try:
            # 首先注册几个设备
            await self.test_device_registration(2)

            # 通过 REST API 发送广播消息
            broadcast_data = {
                "message": "系统维护通知",
                "type": "maintenance",
                "start_time": datetime.now().isoformat(),
                "duration": "30分钟"
            }

            response = requests.post(
                f"{self.base_url}/broadcast",
                json=broadcast_data,
                timeout=10
            )

            if response.status_code == 200:
                result = response.json()
                logger.info(f"✅ 广播消息发送成功 - 成功发送: {result.get('success_count', 0)} 设备")
                return True
            else:
                logger.error(f"❌ 广播消息发送失败: {response.status_code} - {response.text}")
                return False

        except Exception as e:
            logger.error(f"❌ 广播功能测试异常: {str(e)}")
            return False

    async def test_error_handling(self):
        """测试错误处理 - 修复版本"""
        logger.info("🧪 测试错误处理")

        try:
            # 测试发送无效消息
            async with websockets.connect(f"{self.ws_url}/ws/device/error_test") as websocket:
                invalid_message = {
                    "type": "invalid_type",  # 无效的消息类型
                    "data": {"test": "data"},
                    "timestamp": datetime.now().isoformat()
                }

                await websocket.send(json.dumps(invalid_message))
                logger.info("📤 已发送无效消息测试")

                # 应该收到错误响应
                response = await asyncio.wait_for(websocket.recv(), timeout=self.timeout)
                response_data = json.loads(response)

                if response_data.get("type") == "error":
                    logger.info("✅ 错误处理测试成功 - 正确返回错误响应")
                    await websocket.close()
                    return True
                else:
                    logger.error(f"❌ 错误处理测试失败: {response_data}")
                    await websocket.close()
                    return False

        except asyncio.TimeoutError:
            logger.error("⏰ 错误处理测试超时")
            return False
        except Exception as e:
            logger.error(f"❌ 错误处理测试异常: {str(e)}")
            return False

    async def run_comprehensive_test(self):
        """运行全面测试 - 修复版本"""
        logger.info("🚀 开始全面测试 WebSocket 服务器")
        logger.info("=" * 50)

        test_results = {}

        # 1. 测试设备注册
        test_results["device_registration"] = await self.test_device_registration(2)
        await asyncio.sleep(1)

        # 2. 测试心跳功能
        test_results["heartbeat"] = await self.test_heartbeat()
        await asyncio.sleep(1)

        # 3. 测试状态更新
        test_results["status_update"] = await self.test_status_update()
        await asyncio.sleep(1)

        # 4. 测试数据报告
        test_results["data_report"] = await self.test_data_report()
        await asyncio.sleep(1)

        # 5. 测试 REST API
        test_results["rest_api"] = self.test_rest_api()
        await asyncio.sleep(1)

        # 6. 测试广播功能
        test_results["broadcast"] = await self.test_broadcast_functionality()
        await asyncio.sleep(1)

        # 7. 测试错误处理
        test_results["error_handling"] = await self.test_error_handling()
        await asyncio.sleep(1)

        # 8. 测试命令接收（这个测试时间较长，放在最后）
        logger.info("⏳ 开始命令接收测试（可能需要等待30秒）...")
        test_results["command_receiving"] = await self.test_command_receiving()

        # 汇总结果
        logger.info("\n" + "=" * 50)
        logger.info("📊 测试结果汇总")
        logger.info("=" * 50)

        for test_name, result in test_results.items():
            status = "✅ 通过" if result else "❌ 失败"
            logger.info(f"{test_name:20} {status}")

        total_tests = len(test_results)
        passed_tests = sum(1 for result in test_results.values() if result)

        logger.info("=" * 50)
        logger.info(f"总测试数: {total_tests}, 通过: {passed_tests}, 失败: {total_tests - passed_tests}")

        if passed_tests == total_tests:
            logger.info("🎉 所有测试通过！")
        else:
            logger.warning("⚠️  部分测试失败，请检查服务器配置")

        return test_results

    async def cleanup(self):
        """清理测试资源"""
        logger.info("🧹 清理测试资源")
        for device_id, websocket in self.connected_devices.items():
            try:
                await websocket.close()
                logger.info(f"已关闭设备 {device_id} 的连接")
            except:
                pass
        self.connected_devices.clear()


# 快速测试函数
async def quick_test():
    """快速测试函数"""
    tester = WebSocketServerTester()

    try:
        # 只测试关键功能
        print("🚀 快速测试 WebSocket 服务器关键功能")

        # 测试设备注册
        if await tester.test_device_registration(1):
            print("✅ 设备注册测试通过")
        else:
            print("❌ 设备注册测试失败")
            return

        # 测试心跳
        if await tester.test_heartbeat():
            print("✅ 心跳测试通过")
        else:
            print("❌ 心跳测试失败")

        # 测试 REST API
        if tester.test_rest_api():
            print("✅ REST API 测试通过")
        else:
            print("❌ REST API 测试失败")

    finally:
        await tester.cleanup()


# 主测试函数
async def main():
    """主测试函数"""
    tester = WebSocketServerTester()

    try:
        # 运行全面测试
        await tester.run_comprehensive_test()

    except Exception as e:
        logger.error(f"测试过程中发生错误: {e}")
    finally:
        await tester.cleanup()


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "quick":
        # 快速测试模式
        asyncio.run(quick_test())
    else:
        # 全面测试模式
        asyncio.run(main())